package top.sqdpt.rag.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.query.transformer.QueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import top.sqdpt.rag.assistant.RagAssistant;
import top.sqdpt.rag.handler.KnowledgeGraphRetrievalHandler;
import top.sqdpt.rag.store.MysqlChatStore;
import top.sqdpt.rag.tools.AiComputeTool;
import top.sqdpt.rag.tools.WeatherTool;
import top.sqdpt.rag.transformer.CompressingAndExpandingQueryTransformer;

@Configuration
public class AiConfig {

    @Autowired
    private WeatherTool weatherTool;

    @Bean
    public ChatMemoryProvider chatMemoryProvider(MysqlChatStore mysqlChatStore, @Value("${langchain4j.message.max:10}") int max) {
        return memoryId -> MessageWindowChatMemory
                .builder()
                .id(memoryId)
                .maxMessages(max)
                .chatMemoryStore(mysqlChatStore)
                .build();
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(
            @Value("${langchain4j.pg.host}") String host,
            @Value("${langchain4j.pg.port}") Integer port,
            @Value("${langchain4j.pg.username}") String username,
            @Value("${langchain4j.pg.password}") String password,
            @Value("${langchain4j.pg.database}") String database,
            EmbeddingModel embeddingModel) {
        return PgVectorEmbeddingStore.builder()
                .host(host)
                .port(port)
                .user(username)
                .password(password)
                .database(database)
                .table("vx_chat_embeddings")
                .dimension(embeddingModel.dimension())      // Required: Must match the embedding model’s output dimension
                .useIndex(true)                             // Enable IVFFlat index
                .indexListSize(100)                         // Number of lists for IVFFlat index
                .createTable(true)                          // Automatically create the table if it doesn’t exist
                .dropTableFirst(false)                      // Don’t drop the table first (set to true if you want a fresh start)
                .build();
    }

    @Bean
    public RagAssistant aiBidAssistant(ChatLanguageModel chatLanguageModel
            , KnowledgeGraphRetrievalHandler knowledgeGraphRetrievalHandler
            , ChatMemoryProvider chatMemoryProvider) {
        // 检索转换器
        QueryTransformer queryTransformer = new CompressingAndExpandingQueryTransformer(chatLanguageModel);
        // 处理文件时的RAG增强路径
        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .queryTransformer(queryTransformer)
                .contentRetriever(knowledgeGraphRetrievalHandler::retrieveHybrid)
                .build();
        return AiServices.builder(RagAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .retrievalAugmentor(retrievalAugmentor)
                .chatMemoryProvider(chatMemoryProvider)
                .tools(new AiComputeTool(), weatherTool)
                .build();
    }

}
